!! Copyright (C) 2009,2010,2011,2012  Marco Restelli
!!
!! This file is part of:
!!   FEMilaro -- Finite Element Method toolkit
!!
!! FEMilaro is free software; you can redistribute it and/or modify it
!! under the terms of the GNU General Public License as published by
!! the Free Software Foundation; either version 3 of the License, or
!! (at your option) any later version.
!!
!! FEMilaro is distributed in the hope that it will be useful, but
!! WITHOUT ANY WARRANTY; without even the implied warranty of
!! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
!! General Public License for more details.
!!
!! You should have received a copy of the GNU General Public License
!! along with FEMilaro; If not, see <http://www.gnu.org/licenses/>.
!!
!! author: Michele Nini                   <michelenini88@gmail.com>

!<---------------------------------------------------------------------
module mod_smagorinsky_flux

!----------------------------------------------------------------------

 use mod_kinds, only: &
   mod_kinds_initialized, &
   wp

 use mod_messages, only: &
   mod_messages_initialized, &
   error, warning, info

 use mod_grid, only: &
   mod_grid_initialized, &
   t_grid, el_linear_size

 use mod_base, only: &
   mod_base_initialized, &
   t_base

 use mod_sympoly, only: &
   mod_sympoly_initialized, &
   nll_sympoly

 use mod_atm_refstate, only: &
   mod_atm_refstate_initialized, &
   t_atm_refstate
 
 use mod_turb_flux_base, only: &
   mod_turb_flux_base_initialized, &
   c_turbmod, c_turbmod_progs, c_turbmod_diags, &
   t_turb_diags

 use mod_dgcomp_testcases, only: &
   mod_dgcomp_testcases_initialized, &
   t_phc, phc, coeff_visc
 
 use mod_viscous_flux, only: &
   mod_viscous_flux_initialized, &
   t_viscous_flux, compute_vf_flux

!----------------------------------------------------------------------
 
 implicit none

!----------------------------------------------------------------------

! Module interface

 public :: &
   mod_smagorinsky_flux_constructor, &
   mod_smagorinsky_flux_destructor,  &
   mod_smagorinsky_flux_initialized, &
   t_smagorinsky_flux

 private

!----------------------------------------------------------------------

! Module types and parameters

 ! public members

 !> Smagorinsky flux
 !!
 !! The structure resembles that of the viscous flux, with the
 !! addition of two fields used to compute the equivalent element
 !! linear size \f$\Delta\f$ and the element averages.
 !!
 !! The field \c grad stores the gradients of the following quatities,
 !! as for the viscous flux.
 !! <ul>
 !!  <li> temperature (second index equal to 1)
 !!  <li> velocity    (second index from 2 to <tt>1+grid\%d</tt>)
 !!  <li> tracers     (second index \f$\geq\f$<tt>2+grid\%d</tt>)
 !! </ul>
 type, extends(t_viscous_flux) :: t_smagorinsky_flux
  real(wp), allocatable :: ave(:) !< normalized weights (averages)
  real(wp), allocatable :: delta2(:) !< squared element linear size
 contains 
  procedure, pass(tm) :: init                => sf_init
  procedure, pass(tm) :: clean               => sf_clean
  ! The gradient is computed as in the viscous case: no need to
  ! redefine compute_grad_diags.
  procedure, pass(tm) :: compute_coeff_diags => sf_compute_coeff_diags
  procedure, pass(tm) :: flux                => sf_flux
 end type t_smagorinsky_flux

 ! No need to extend c_turbmod_progs

 !> Smagorinsky diagnostics.
 !!
 !! The base type \c c_turbmod_diags must be extended to include the
 !! element averaged deformation rate.
 type, extends(c_turbmod_diags) :: t_smagorinsky_diags
  real(wp), allocatable :: s2(:) !< \f$\|\nabla^S\underline{u}\|^2\f$
 end type t_smagorinsky_diags

! Module variables
 ! public members
 logical, protected ::               &
   mod_smagorinsky_flux_initialized = .false.
 character(len=*), parameter :: &
   this_mod_name = 'mod_smagorinsky_flux'

!----------------------------------------------------------------------

contains

!----------------------------------------------------------------------

 subroutine mod_smagorinsky_flux_constructor()
  character(len=*), parameter :: &
   this_sub_name = 'constructor'

   !Consistency checks ---------------------------
   if( (mod_kinds_initialized.eqv..false.) .or. &
     (mod_messages_initialized.eqv..false.) .or. &
     (mod_grid_initialized.eqv..false.) .or. &
     (mod_base_initialized.eqv..false.) .or. &
     (mod_atm_refstate_initialized.eqv..false.) .or. &
     (mod_turb_flux_base_initialized.eqv..false.).or.&
     (mod_dgcomp_testcases_initialized.eqv..false.).or.&
     (mod_viscous_flux_initialized.eqv..false.)) then
     call error(this_sub_name,this_mod_name, &
                'Not all the required modules are initialized.')
   endif
   if(mod_smagorinsky_flux_initialized.eqv..true.) then
     call warning(this_sub_name,this_mod_name, &
                  'Module is already initialized.')
   endif

   mod_smagorinsky_flux_initialized = .true.

 end subroutine mod_smagorinsky_flux_constructor

!----------------------------------------------------------------------

 subroutine mod_smagorinsky_flux_destructor()
  character(len=*), parameter :: &
    this_sub_name = 'destructor'
   
   !Consistency checks ---------------------------
   if(mod_smagorinsky_flux_initialized.eqv..false.) then
     call error(this_sub_name,this_mod_name, &
                'This module is not initialized.')
   endif
   !----------------------------------------------

   mod_smagorinsky_flux_initialized = .false.
 end subroutine mod_smagorinsky_flux_destructor

!----------------------------------------------------------------------

 pure subroutine sf_init(tm,progs,diags,td,grid,base,ntrcs)
  type(t_grid),       intent(in) :: grid
  type(t_base),       intent(in) :: base
  integer,            intent(in) :: ntrcs
  class(t_smagorinsky_flux), intent(inout) :: tm
  class(c_turbmod_progs), allocatable, intent(out) :: progs
  class(c_turbmod_diags), allocatable, intent(out) :: diags
  type(t_turb_diags), intent(out) :: td

  integer :: ie
  real(wp) :: ngdl ! number of dofs, as a real

   ! model
   tm%d           = grid%d
   tm%ntrcs       = ntrcs
   tm%initialized = .true.
   allocate(tm%wg(base%m));             tm%wg     = base%wg
   allocate(tm%base_p(base%pk,base%m)); tm%base_p = base%p
   ! specific fields of the Smagorinsky model
   allocate(tm%ave(base%m));            tm%ave    = base%wg/base%me%vol
   ! element significant length
   allocate(tm%delta2(grid%ne))
   do ie=1,grid%ne
     ngdl = real(nll_sympoly(tm%d,base%k),wp)**(1.0_wp/real(tm%d,wp))
     tm%delta2(ie) = (el_linear_size(grid%e(ie)) / ngdl)**2
   enddo

   ! no need to allocate progs

   ! diags
   allocate(t_smagorinsky_diags::diags)
   diags%ngrad    = 1 + tm%d + tm%ntrcs
   allocate(diags%grad(tm%d,diags%ngrad,base%pk,grid%ne))
   select type(diags); type is(t_smagorinsky_diags)
    allocate(diags%s2(grid%ne))
   end select

   ! diagnostics
   td%ndiags = 2
   allocate( td%diag_names(td%ndiags) ,  &
     td%diags(td%ndiags,base%pk,grid%ne) )
   td%diag_names(1) = '|nabla^S u + lambda*div(u)|'
   td%diag_names(2) = 'dynamic dissipation'

 end subroutine sf_init

!----------------------------------------------------------------------

 pure subroutine sf_clean(tm,diags,td)
  class(t_smagorinsky_flux), intent(inout) :: tm
  class(c_turbmod_diags), allocatable, intent(inout) :: diags
  type(t_turb_diags), intent(out) :: td
   
   deallocate(tm%wg,tm%base_p)
   deallocate(tm%ave,tm%delta2)
   deallocate(diags) ! deallocates all the components
   ! td components deallocated automatically

 end subroutine sf_clean

!-----------------------------------------------------------------------

 pure subroutine sf_compute_coeff_diags(tm,cd , grid,base, &
                                        uuu,progs,atm_ref)
  class(t_smagorinsky_flux), intent(in) :: tm
  type(t_grid),           intent(in)    :: grid
  type(t_base),           intent(in)    :: base   
  real(wp),               intent(in)    :: uuu(:,:,:)
  class(t_atm_refstate),  intent(in)    :: atm_ref(:,:)
  class(c_turbmod_progs), allocatable, intent(in) :: progs
  class(c_turbmod_diags), intent(inout) :: cd

  ! local variables
  integer  :: d, ie, id, jd, l, i
  real(wp) :: s2(base%m)
  real(wp) :: gradg(grid%d,grid%d,base%m), gsuu(grid%d,grid%d,base%m)
  
   select type(cd); type is(t_smagorinsky_diags)

   d = grid%d
   do ie=1,grid%ne

     ! compute the symmetric velocity gradient
     gradg = 0.0_wp 
     do l= 1,size(base%p,2)
       do i=1,size(base%p,1)
         gradg(:,:,l) = gradg(:,:,l) + cd%grad(:,2:1+d,i,ie)*base%p(i,l)
       enddo
     enddo
     do id=1,d
       do jd=1,id-1; gsuu(id,jd,:) = gsuu(jd,id,:); enddo ! symmetry
       do jd=id,d
         gsuu(id,jd,:) = gradg(jd,id,:)+gradg(id,jd,:)
       enddo
     enddo

     ! local strain rate
     do l=1,base%m
       s2(l) = sum( gsuu(:,:,l)**2 )
     enddo

     ! element average
     cd%s2(ie) = sum( tm%ave*s2 )
        
   enddo

   end select
 end subroutine sf_compute_coeff_diags

!----------------------------------------------------------------------

 pure subroutine sf_flux(fem, tm,ie,x,bp, consv,atm_ref, rho,p,uu,cc, &
                         progs,diags, td)
  class(t_smagorinsky_flux), intent(in) :: tm
  integer,               intent(in) :: ie
  real(wp),              intent(in) :: x(:,:)
  real(wp),              intent(in) :: bp(:,:)
  real(wp),              intent(in) :: consv(:,:)
  class(t_atm_refstate), intent(in) :: atm_ref(:)
  real(wp),              intent(in) :: rho(:), p(:), uu(:,:), cc(:,:)
  class(c_turbmod_progs), allocatable, intent(in) :: progs
  class(c_turbmod_diags), intent(in) :: diags
  real(wp),              intent(out) :: fem(:,:,:)
  type(t_turb_diags),    intent(inout), optional :: td

  ! Model parameter
  real(wp), parameter :: &
    cs     = 0.1_wp, &
    ci     = 0.0_wp, &
    pr_sgs = 0.72_wp 

  ! local variables 
  integer  :: l, id
  real(wp) :: nu_sgs(size(x,2)), tau_kk, sym(tm%d,tm%d,size(x,2)), &
    tau_sgs(tm%d,tm%d), gradg(tm%d,diags%ngrad,size(x,2))
                          
   ! viscous contribution      
   call compute_vf_flux(fem, tm,ie,x,bp, consv,atm_ref, rho,p,uu,cc, &
                        diags, td=td, vf_sym=sym , vf_gradg=gradg)

   select type(diags); type is(t_smagorinsky_diags)
    
   ! add the turbulent contribution
   do l=1,size(x,2)

     nu_sgs(l) = cs**2 * tm%delta2(ie) * sqrt(diags%s2(ie)/2.0_wp)     &
             ! van Driest damping
             *( 1.0_wp - exp(-(1.0_wp-abs(x(2,l)))*phc%re_tau/25.0_wp) )
                     
     ! momentum flux
     ! deviatoric sgs stress tensor
     tau_sgs = - rho(l)*nu_sgs(l)*sym(:,:,l)
     ! isotropic sgs stress tensor
     tau_kk = rho(l)*ci*tm%delta2(ie) * diags%s2(ie)/2.0_wp 
   
     do id=1,tm%d
       ! check whether 3.0 should be d
       tau_sgs(id,id) = tau_sgs(id,id) - 1.0_wp/3.0_wp*tau_kk
     enddo
     fem(:,2:1+tm%d,l) = fem(:,2:1+tm%d,l) + tau_sgs 
 
     ! energy flux
     fem(:,   1    ,l) = fem(:,   1    ,l)   &
     ! sgs heat flux
     - rho(l)*nu_sgs(l)*phc%cp/pr_sgs*gradg(:,1,l) & 
     ! turbulent diffusion (1): Knight98 -> uk*tau_jk^sgs
     ! + 0.5*uj*tau_kk^sgs
     + matmul( uu(:,l), transpose(tau_sgs))   &
     + 0.5_wp*uu(:,l)*tr( tau_sgs )

     if (present(td)) then       
       ! todo: implement this, combining it with the diagnostics of
       ! the viscous flux.
     endif

   enddo

   end select
 end subroutine sf_flux
 
!----------------------------------------------------------------------

 pure function tr(a)
  real(wp), intent(in) :: a(:,:)
  real(wp) :: tr
  integer :: i

   tr = a(1,1)
   do i=2,minval(shape(a))
     tr = tr + a(i,i)
   enddo
 end function tr

!-----------------------------------------------------------------------

end module mod_smagorinsky_flux

